import torch
import torch.nn as nn


class EasyQNet(nn.Module):
    NEED_COL = ["bbi_signal", "bias_signal", "cci_signal",
                "dma_signal", "expma_signal", "kdj_signal", "macd_signal",
                "mfi_signal", "mi_signal", "mtm_signal", "priceosc_signal",
                "psy_signal", "roc_signal"]

    def __init__(self, state_width=len(NEED_COL), action_size=2, alpha=0.01, dropout=0.2):
        super(EasyQNet, self).__init__()
        self.linear = nn.Linear(state_width, action_size)
        self.alpha = alpha

    def forward(self, x, position):
        out = self.linear(x)
        return out

    def l1_penalty(self):
        l1_norm = torch.norm(self.linear.weight, p=1)
        return self.alpha * l1_norm
